import os
import shutil
import cv2
import xml.etree.ElementTree as ET


##get object annotation bndbox loc start
def GetAnnotBoxLoc(AnotPath):  # AnotPath VOC标注文件路径
    tree = ET.ElementTree(file=AnotPath)  # 打开文件，解析成一棵树型结构
    root = tree.getroot()  # 获取树型结构的根
    ObjectSet = root.findall('object')  # 找到文件中所有含有object关键字的地方，这些地方含有标注目标
    ObjBndBoxSet = {}  # 以目标类别为关键字，目标框为值组成的字典结构
    for Object in ObjectSet:
        ObjName = Object.find('name').text
        BndBox = Object.find('bndbox')
        x1 = int(BndBox.find('xmin').text)  # -1 #-1是因为程序是按0作为起始位置的
        y1 = int(BndBox.find('ymin').text)  # -1
        x2 = int(BndBox.find('xmax').text)  # -1
        y2 = int(BndBox.find('ymax').text)  # -1
        BndBoxLoc = [x1, y1, x2, y2]
        if ObjName in ObjBndBoxSet:
            ObjBndBoxSet[ObjName].append(BndBoxLoc)  # 如果字典结构中含有这个类别了，那么这个目标框要追加到其值的末尾
        else:
            ObjBndBoxSet[ObjName] = [BndBoxLoc]  # 如果字典结构中没有这个类别，那么这个目标框就直接赋值给其值吧
    return ObjBndBoxSet


##get object annotation bndbox loc end
def display(objBox, pic, file_name):
    img = pic

    for key in objBox.keys():
        for i in range(len(objBox[key])):
            cv2.rectangle(img, (objBox[key][i][0], objBox[key][i][1]), (objBox[key][i][2], objBox[key][i][3]),
                          (0, 0, 255), 2)
            cv2.putText(img, key, (objBox[key][i][0], objBox[key][i][1]), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 1)
    # cv2.imshow('img', img)
    cv2.imwrite(file_name, img)
    # cv2.waitKey(0)

def change_rate_box(objBox, w_rate, h_rate):
    new_objBox = {}
    for key in objBox.keys():
        for i in range(len(objBox[key])):
            x1 = int(objBox[key][i][0] * w_rate)
            y1 = int(objBox[key][i][1] * h_rate)
            x2 = int(objBox[key][i][2] * w_rate)
            y2 = int(objBox[key][i][3] * h_rate)
            BndBoxLoc = [x1, y1, x2, y2]
            if key in new_objBox:
                new_objBox[key].append(BndBoxLoc)  # 如果字典结构中含有这个类别了，那么这个目标框要追加到其值的末尾
            else:
                new_objBox[key] = [BndBoxLoc]  # 如果字典结构中没有这个类别，那么这个目标框就直接赋值给其值吧
    return new_objBox

base_root = '/data/ICE/RGBCollection/base'
base_image_path = os.path.join(base_root, 'images')
base_xml_path = os.path.join(base_root, 'xml')
image_names = os.listdir(base_image_path)

rgb_path = '/data/ICE/RGBCollection/rgb_images'
for index_image in image_names[20:21]:
    print(index_image)
    base_image = cv2.imread(os.path.join(base_image_path, index_image))
    rgb_image = cv2.imread(os.path.join(rgb_path, index_image))
    if rgb_image is None:
        continue
    rows, cols, _ = rgb_image.shape
    M = cv2.getRotationMatrix2D(((cols - 1) / 2.0, (rows - 1) / 2.0), 180, 1)
    dst_image = cv2.warpAffine(rgb_image, M, (cols, rows))
    cv2.imwrite('base.jpg', base_image)
    cv2.imwrite('rgb.jpg', dst_image)
    # cv2.waitKey()

    xml_path = os.path.join(base_xml_path, index_image[:-4] + '.xml')
    ObjBndBoxSet = GetAnnotBoxLoc(xml_path)
    new_ObjBndBoxSet = change_rate_box(ObjBndBoxSet, 3, 2.7)
    print(ObjBndBoxSet)
    print(new_ObjBndBoxSet)
    display(ObjBndBoxSet, base_image, 'display_base.jpg')
    display(new_ObjBndBoxSet, dst_image, 'display_rgb.jpg')

    base_image_resize = cv2.resize(base_image, (1920, 1080))
    display(new_ObjBndBoxSet, base_image_resize, 'display_resize_base.jpg')




